import os
import uvicorn
import time
import json
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.responses import JSONResponse, FileResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from typing import List, Optional
from pathlib import Path
import shutil
import logging
from pdf_processor import PDFProcessor
from ai_engine import AIEngine
from config import APP_HOST, APP_PORT, DEBUG, UPLOADS_DIR, DASHSCOPE_API_KEY

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 检查API密钥
if not DASHSCOPE_API_KEY:
    logger.error("未设置DASHSCOPE_API_KEY环境变量，请先设置API密钥")
    exit(1)

# 创建应用
app = FastAPI(title="网络实验AI助手", description="基于通义千问API的网络实验PDF处理和问答系统")

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 初始化组件
pdf_processor = PDFProcessor()
ai_engine = AIEngine(pdf_processor)

# 挂载静态文件
app.mount("/static", StaticFiles(directory="static"), name="static")
app.mount("/uploads", StaticFiles(directory=str(UPLOADS_DIR)), name="uploads")

@app.get("/")
async def root():
    """API根路径，返回系统状态"""
    return {"status": "online", "message": "网络实验AI助手API已启动"}

@app.post("/upload-pdf")
async def upload_pdf(
    file: UploadFile = File(...),
    is_guide: bool = Form(True),
    include_images: bool = Form(True)
):
    """上传PDF文件并处理
    
    Args:
        file: PDF文件
        is_guide: 是否为实验指导书（否则为实验报告）
        include_images: 是否提取并识别PDF中的图片
    """
    if not file.filename.lower().endswith('.pdf'):
        raise HTTPException(status_code=400, detail="只接受PDF文件")
    
    # 保存文件
    file_path = Path(UPLOADS_DIR) / file.filename
    try:
        with open(file_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
    except Exception as e:
        logger.error(f"保存文件失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"保存文件失败: {str(e)}")
    
    # 处理PDF（传递include_images参数）
    success = pdf_processor.process_pdf(str(file_path), is_guide, include_images)
    
    if success:
        return {
            "success": True,
            "message": f"成功处理 {file.filename}" + (" (含图片识别)" if include_images else ""),
            "file_path": str(file_path),
            "type": "指导书" if is_guide else "实验报告",
            "with_images": include_images
        }
    else:
        raise HTTPException(status_code=500, detail=f"处理PDF失败")

@app.post("/batch-upload")
async def batch_upload(
    files: List[UploadFile] = File(...),
    is_guide: bool = Form(True)
):
    """批量上传PDF文件
    
    Args:
        files: PDF文件列表
        is_guide: 是否为实验指导书（否则为实验报告）
    """
    results = []
    
    for file in files:
        if not file.filename.lower().endswith('.pdf'):
            results.append({
                "filename": file.filename,
                "success": False,
                "message": "只接受PDF文件"
            })
            continue
        
        # 保存文件
        file_path = Path(UPLOADS_DIR) / file.filename
        try:
            with open(file_path, "wb") as buffer:
                shutil.copyfileobj(file.file, buffer)
        except Exception as e:
            results.append({
                "filename": file.filename,
                "success": False,
                "message": f"保存文件失败: {str(e)}"
            })
            continue
        
        # 处理PDF
        success = pdf_processor.process_pdf(str(file_path), is_guide)
        
        results.append({
            "filename": file.filename,
            "success": success,
            "message": "处理成功" if success else "处理失败",
            "file_path": str(file_path),
            "type": "指导书" if is_guide else "实验报告"
        })
    
    return {"results": results}

@app.post("/query")
async def query(
    question: str = Form(...),
    search_docs: bool = Form(True),
    conversation_id: str = Form(None)
):
    """文本查询
    
    Args:
        question: 用户问题
        search_docs: 是否搜索文档
        conversation_id: 对话ID，如果提供则加载该对话
    """
    result = ai_engine.text_query(question, search_docs, conversation_id)
    return result

@app.get("/api/image/{image_name:path}")
@app.get("/uploads/{image_name:path}")  # 保留旧路径以兼容
async def get_image(image_name: str):
    """获取上传的图片"""
    try:
        # 直接记录请求的图片名称，用于调试
        logger.info(f"请求图片: {image_name}")
        
        # 尝试直接查找完整文件名
        direct_path = UPLOADS_DIR / image_name
        if direct_path.exists():
            logger.info(f"直接找到文件: {direct_path}")
            return FileResponse(str(direct_path))
        
        # 尝试查找所有图片文件
        uploads_path = Path(UPLOADS_DIR)
        logger.info(f"查找目录: {uploads_path}")
        
        # 确保目录存在
        if uploads_path.exists() and uploads_path.is_dir():
            # 列出目录中的所有文件
            all_files = list(uploads_path.iterdir())
            logger.info(f"目录中的文件数量: {len(all_files)}")
            
            # 首先尝试查找完全匹配的文件名（不区分大小写）
            for file_path in all_files:
                if file_path.is_file():
                    file_name = file_path.name
                    logger.info(f"检查文件: {file_name}")
                    
                    # 检查完全匹配（不区分大小写）
                    if file_name.lower() == image_name.lower():
                        logger.info(f"找到完全匹配文件: {file_path}")
                        return FileResponse(str(file_path))
            
            # 如果没有完全匹配，尝试查找部分匹配
            # 提取前缀（时间戳部分）
            parts = image_name.split('_')
            if len(parts) > 0:
                prefix = parts[0]
                logger.info(f"尝试查找前缀为 {prefix} 的文件")
                
                for file_path in all_files:
                    if file_path.is_file():
                        file_name = file_path.name
                        if prefix in file_name:
                            logger.info(f"找到前缀匹配文件: {file_path}")
                            return FileResponse(str(file_path))
            
            # 如果前缀匹配也失败，尝试查找任何图片文件
            logger.info("尝试查找任何图片文件")
            for file_path in all_files:
                if file_path.is_file() and file_path.suffix.lower() in ['.jpg', '.jpeg', '.png']:
                    logger.info(f"找到图片文件: {file_path}")
                    return FileResponse(str(file_path))
        
        # 如果文件不存在，返回404错误
        logger.error(f"图片不存在: {image_name}")
        return JSONResponse(
            status_code=404,
            content={"error": "图片不存在"}
        )
    except Exception as e:
        logger.error(f"获取图片时出错: {str(e)}")
        return JSONResponse(
            status_code=500,
            content={"error": f"获取图片时出错: {str(e)}"}
        )

@app.post("/multimodal-query")
async def multimodal_query(
    question: str = Form(...),
    images: List[UploadFile] = File(...),
    conversation_id: str = Form(None)
):
    """多模态查询（文本+图像）
    
    Args:
        question: 用户问题
        images: 图像文件列表
        conversation_id: 对话ID，如果提供则加载该对话
    """
    image_paths = []
    
    # 记录请求信息
    logger.info(f"收到多模态查询请求: 问题='{question}', 图片数量={len(images)}")
    
    # 保存上传的图像
    if images:
        for i, img in enumerate(images):
            if not img.filename:
                logger.warning(f"图片 {i+1} 没有文件名，跳过")
                continue
                
            # 检查文件类型
            supported_formats = ['.jpg', '.jpeg', '.png', '.gif', '.webp', '.bmp', '.tiff', '.tif']
            if not any(img.filename.lower().endswith(ext) for ext in supported_formats):
                logger.warning(f"图片 {img.filename} 不是支持的格式，跳过。支持的格式: {', '.join(supported_formats)}")
                continue
            
            # 生成唯一文件名，避免覆盖
            timestamp = int(time.time() * 1000)  # 毫秒级时间戳
            unique_filename = f"{timestamp}_{i}_{img.filename}"
            img_path = Path(UPLOADS_DIR) / unique_filename
            
            try:
                # 读取文件内容
                content = await img.read()
                
                # 检查文件大小
                file_size = len(content)
                logger.info(f"图片 {img.filename} 大小: {file_size} 字节")
                
                if file_size == 0:
                    logger.warning(f"图片 {img.filename} 为空文件，跳过")
                    continue
                    
                # 检查文件大小限制（10MB）
                if file_size > 10 * 1024 * 1024:
                    logger.warning(f"图片 {img.filename} 大小超过10MB限制，跳过")
                    continue
                
                # 保存文件
                with open(img_path, "wb") as buffer:
                    buffer.write(content)
                
                # 检查保存的文件是否存在
                if img_path.exists():
                    image_paths.append(str(img_path))
                    logger.info(f"保存图像成功: {img_path}")
                else:
                    logger.error(f"保存后文件不存在: {img_path}")
            except Exception as e:
                logger.error(f"保存图像 {img.filename} 失败: {str(e)}")
    
    if not image_paths:
        logger.warning("没有有效的图像文件可供处理")
        return {
            "success": False,
            "error": "没有有效的图像文件"
        }
    
    # 调用多模态查询
    logger.info(f"开始处理多模态查询，有效图片数量: {len(image_paths)}")
    result = ai_engine.multimodal_query(question, image_paths, conversation_id)
    logger.info(f"多模态查询处理完成，结果: {result['success']}")
    return result

@app.post("/add-memory")
async def add_memory(
    fact: str = Form(...),
    is_correction: bool = Form(False)
):
    """添加记忆
    
    Args:
        fact: 要记住的事实
        is_correction: 是否为纠正信息
    """
    ai_engine.add_to_memory(fact, is_correction)
    return {
        "success": True,
        "message": "成功添加到记忆",
        "is_correction": is_correction
    }

@app.get("/memory")
async def get_memory():
    """获取当前记忆"""
    return ai_engine.memory

@app.get("/conversations")
@app.get("/api/conversations")  # 保留两个路径以兼容
async def get_conversations():
    """获取所有对话列表"""
    try:
        conversations = ai_engine.get_conversations_list()
        return {"success": True, "conversations": conversations}
    except Exception as e:
        logger.error(f"获取对话列表失败: {str(e)}")
        return {"success": False, "error": str(e), "conversations": []}

@app.post("/create-conversation")
@app.post("/api/create-conversation")  # 保留两个路径以兼容
async def create_conversation():
    """创建新的对话"""
    conversation_id = ai_engine.create_new_conversation()
    return {"success": True, "conversation_id": conversation_id}

@app.post("/load-conversation")
@app.post("/api/load-conversation")  # 保留两个路径以兼容
async def load_conversation(conversation_id: str = Form(...)):
    """加载指定的对话"""
    success = ai_engine.load_conversation(conversation_id)
    
    # 如果成功加载，返回对话历史
    if success:
        # 获取对话文件路径
        conversation_file = ai_engine.conversations_dir / f"{conversation_id}.json"
        
        try:
            # 读取对话内容
            with open(conversation_file, 'r', encoding='utf-8') as f:
                conversation_data = json.load(f)
                
            # 返回对话历史
            return {
                "success": True, 
                "conversation_id": conversation_id,
                "title": conversation_data.get("title", ""),
                "messages": conversation_data.get("messages", [])
            }
        except Exception as e:
            logger.error(f"读取对话内容失败: {str(e)}")
            return {"success": False, "error": f"读取对话内容失败: {str(e)}"}
    else:
        return {"success": False, "conversation_id": None, "error": "加载对话失败"}

@app.post("/delete-conversation")
@app.post("/api/delete-conversation")  # 保留两个路径以兼容
async def delete_conversation(conversation_id: str = Form(...)):
    """删除指定的对话"""
    try:
        import os
        from pathlib import Path
        
        # 构建对话文件路径
        conversation_file = ai_engine.conversations_dir / f"{conversation_id}.json"
        
        # 检查文件是否存在
        if not conversation_file.exists():
            return {"success": False, "error": "对话不存在"}
        
        # 删除文件
        os.remove(conversation_file)
        
        # 如果删除的是当前对话，重置当前对话ID
        if ai_engine.current_conversation_id == conversation_id:
            ai_engine.current_conversation_id = None
            ai_engine.conversation_history = []
        
        return {"success": True}
    except Exception as e:
        logger.error(f"删除对话失败: {str(e)}")
        return {"success": False, "error": str(e)}

@app.post("/delete-memory")
async def delete_memory(
    index: int = Form(...),
    is_correction: bool = Form(False)
):
    """删除记忆
    
    Args:
        index: 要删除的记忆索引
        is_correction: 是否为纠正信息
    """
    success = ai_engine.delete_from_memory(index, is_correction)
    return {
        "success": success,
        "message": "成功删除记忆" if success else "删除记忆失败，索引无效",
        "is_correction": is_correction
    }

@app.post("/direct-search")
async def direct_search(
    query: str = Form(...),
    collection_name: str = Form("both"),
    limit: int = Form(5),
    search_type: str = Form("semantic")
):
    """直接检索文档
    
    Args:
        query: 查询文本
        collection_name: 要搜索的集合，可以是"guides"、"reports"或"both"
        limit: 返回结果数量限制
        search_type: 搜索类型，"semantic"为语义搜索，"exact"为精确文本匹配
    """
    if search_type == "exact":
        results = pdf_processor.search_exact_text(query, collection_name, limit)
    else:
        results = pdf_processor.search_similar(query, collection_name, limit)
    return {"results": results, "search_type": search_type}

@app.get("/list-pdfs")
async def list_pdfs():
    """列出所有上传的PDF文件"""
    guides = []
    reports = []
    
    try:
        # 获取指导书集合中的所有文档
        guide_docs = pdf_processor.guide_collection.get()
        if guide_docs and guide_docs["metadatas"]:
            # 提取唯一的文件名
            filenames = set()
            for metadata in guide_docs["metadatas"]:
                if "filename" in metadata:
                    filenames.add(metadata["filename"])
            guides = list(filenames)
        
        # 获取报告集合中的所有文档
        report_docs = pdf_processor.report_collection.get()
        if report_docs and report_docs["metadatas"]:
            # 提取唯一的文件名
            filenames = set()
            for metadata in report_docs["metadatas"]:
                if "filename" in metadata:
                    filenames.add(metadata["filename"])
            reports = list(filenames)
        
        return {
            "success": True,
            "guides": guides,
            "reports": reports
        }
    except Exception as e:
        logger.error(f"获取PDF列表时出错: {str(e)}")
        return {
            "success": False,
            "error": str(e),
            "guides": [],
            "reports": []
        }

@app.post("/delete-pdf")
async def delete_pdf(
    filename: str = Form(...),
    is_guide: bool = Form(True)
):
    """删除PDF文件及其向量数据库记录
    
    Args:
        filename: PDF文件名
        is_guide: 是否为实验指导书（否则为实验报告）
    """
    try:
        # 删除向量数据库中的记录
        db_success = pdf_processor.delete_pdf(filename, is_guide)
        
        # 尝试删除物理文件
        file_path = Path(UPLOADS_DIR) / filename
        file_deleted = False
        if file_path.exists():
            try:
                file_path.unlink()
                file_deleted = True
                logger.info(f"已删除文件: {file_path}")
            except Exception as e:
                logger.warning(f"删除文件失败: {str(e)}")
        
        return {
            "success": db_success,
            "message": f"成功删除: {filename}" if db_success else f"删除失败: {filename}",
            "db_deleted": db_success,
            "file_deleted": file_deleted
        }
    except Exception as e:
        logger.error(f"删除PDF时出错: {str(e)}")
        raise HTTPException(status_code=500, detail=f"删除失败: {str(e)}")
        
@app.get("/api/list-images", response_model=dict)
@app.get("/list-images", response_model=dict)  # 保留旧路径以兼容
async def list_images():
    """列出所有上传的图片文件"""
    images = []
    
    try:
        # 获取上传目录中的所有图片文件
        uploads_path = Path(UPLOADS_DIR)
        image_extensions = ['.jpg', '.jpeg', '.png', '.gif', '.webp', '.bmp', '.tiff', '.tif']
        
        # 确保上传目录存在
        uploads_path.mkdir(exist_ok=True, parents=True)
        
        for file_path in uploads_path.iterdir():
            if file_path.is_file() and any(file_path.suffix.lower() == ext for ext in image_extensions):
                # 获取文件信息
                file_stat = file_path.stat()
                
                # 使用文件名作为图片路径
                file_name = file_path.name
                # 使用直接的静态文件路径
                relative_path = f"/uploads/{file_name}"
                
                # 记录图片信息，用于调试
                logger.info(f"添加图片到列表: {file_name}, 路径: {relative_path}")
                
                images.append({
                    "filename": file_name,
                    "path": relative_path,  # 使用新的图片API路径
                    "size": file_stat.st_size,
                    "modified": file_stat.st_mtime
                })
        
        # 按修改时间排序，最新的在前面
        images.sort(key=lambda x: x["modified"], reverse=True)
        
        return {
            "success": True,
            "images": images
        }
    except Exception as e:
        logger.error(f"获取图片列表时出错: {str(e)}")
        return {
            "success": False,
            "error": str(e),
            "images": []
        }

@app.post("/get-pdf-content")
async def get_pdf_content(
    filename: str = Form(...),
    is_guide: bool = Form(True)
):
    """获取PDF文件内容
    
    Args:
        filename: 文件名
        is_guide: 是否为实验指导书（否则为实验报告）
    """
    try:
        # 选择合适的集合
        collection = pdf_processor.guide_collection if is_guide else pdf_processor.report_collection
        
        # 查询匹配文件名的所有页面
        query_result = collection.get(
            where={"filename": filename}
        )
        
        if not query_result or not query_result["documents"]:
            return {
                "success": False,
                "error": "未找到指定文件",
                "pages": []
            }
        
        # 整理页面内容
        pages = []
        for i, (doc, metadata) in enumerate(zip(query_result["documents"], query_result["metadatas"])):
            pages.append({
                "page_number": metadata.get("page", i + 1),
                "content": doc
            })
        
        # 按页码排序
        pages.sort(key=lambda x: x["page_number"])
        
        return {
            "success": True,
            "filename": filename,
            "type": "指导书" if is_guide else "报告",
            "pages": pages
        }
    except Exception as e:
        logger.error(f"获取PDF内容时出错: {str(e)}")
        return {
            "success": False,
            "error": str(e),
            "pages": []
        }

if __name__ == "__main__":
    uvicorn.run("app:app", host=APP_HOST, port=APP_PORT, reload=DEBUG)
